#!/bin/bash

# Swap大小调整脚本（优化依赖版本）
# 优先使用awk进行数学运算，bc作为可选依赖

set -e  # 遇到错误立即退出

# 颜色定义
RED='\033[0;31m'
GREEN='\033[0;32m'
YELLOW='\033[1;33m'
BLUE='\033[0;34m'
CYAN='\033[0;36m'
NC='\033[0m' # No Color

# 全局变量
SWAP_FILE_PATH="/swapfile"
CURRENT_SWAP_MB=0
CURRENT_SWAP_GB=0
SWAP_SIZE_MB=0
SWAP_SIZE_GB=0
HAS_BC=0

# 打印颜色输出函数
print_color() {
    echo -e "${1}${2}${NC}"
}

# 检查命令依赖
check_dependencies() {
    print_color $BLUE "=== 检查系统依赖 ==="
    
    # 检查awk (必需)
    if ! command -v awk &> /dev/null; then
        print_color $RED "❌ 错误: awk命令不存在"
        print_color $RED "   awk是系统基本组件，您的系统可能不完整"
        exit 1
    else
        print_color $GREEN "✓ awk命令可用"
    fi
    
    # 检查bc (可选，用于更精确的计算)
    if command -v bc &> /dev/null; then
        HAS_BC=1
        print_color $GREEN "✓ bc命令可用（用于精确计算）"
    else
        HAS_BC=0
        print_color $YELLOW "⚠️  bc命令未安装，使用awk进行数学运算"
        print_color $CYAN "   如需安装bc，可运行: sudo apt-get install bc"
    fi
    echo ""
}

# 计算GB显示（优先使用awk，bc作为备选）
calculate_gb() {
    local mb=$1
    if [[ $HAS_BC -eq 1 ]]; then
        # 使用bc进行精确计算
        echo "scale=1; $mb/1024" | bc
    else
        # 使用awk作为替代方案
        echo "$mb 1024" | awk '{printf "%.1f", $1/$2}'
    fi
}

# 浮点数比较（优先使用awk，bc作为备选）
float_compare() {
    local a=$1
    local op=$2
    local b=$3
    
    if [[ $HAS_BC -eq 1 ]]; then
        # 使用bc进行精确比较
        echo "$a $op $b" | bc -l
    else
        # 使用awk进行浮点数比较
        case $op in
            ">") echo "$a $b" | awk '{exit ($1 > $2) ? 0 : 1}' ;;
            ">=") echo "$a $b" | awk '{exit ($1 >= $2) ? 0 : 1}' ;;
            "<") echo "$a $b" | awk '{exit ($1 < $2) ? 0 : 1}' ;;
            "<=") echo "$a $b" | awk '{exit ($1 <= $2) ? 0 : 1}' ;;
            "==") echo "$a $b" | awk '{exit ($1 == $2) ? 0 : 1}' ;;
            *) return 1 ;;
        esac
    fi
}

# 修复的权限检查函数
check_root() {
    # 检查是否通过sudo运行
    if [[ $EUID -eq 0 ]] && [[ -z "$SUDO_USER" ]]; then
        print_color $RED "错误：检测到直接使用root用户运行"
        print_color $RED "请使用普通用户身份运行，并在需要时输入密码："
        print_color $YELLOW "  $ sudo ./$(basename "$0")"
        exit 1
    fi
    
    # 检查是否没有root权限
    if [[ $EUID -ne 0 ]]; then
        print_color $RED "错误：此脚本需要root权限"
        print_color $YELLOW "请使用sudo运行："
        print_color $YELLOW "  $ sudo ./$(basename "$0")"
        exit 1
    fi
    
    # 如果通过sudo运行且有root权限，则通过检查
    print_color $GREEN "✓ 权限检查通过（通过sudo运行）"
    echo ""
}

# 检查系统兼容性
check_compatibility() {
    print_color $BLUE "=== 系统兼容性检查 ==="
    
    # 检查操作系统
    if [[ -f /etc/os-release ]]; then
        source /etc/os-release
        print_color $GREEN "✓ 操作系统: $NAME $VERSION"
    else
        print_color $YELLOW "⚠️  无法确定操作系统类型，继续执行..."
    fi
    
    # 检查是否使用Swap文件
    local swap_type=$(swapon --show 2>/dev/null | grep -v "NAME" | awk '{print $2}' | head -1)
    if [[ "$swap_type" != "file" ]] && [[ -n "$swap_type" ]]; then
        print_color $RED "❌ 当前系统使用的是Swap分区，不是Swap文件"
        print_color $RED "   本脚本仅支持调整Swap文件，不支持Swap分区"
        exit 1
    fi
    
    # 检查Swap文件路径
    local swap_file=$(swapon --show 2>/dev/null | grep -v "NAME" | awk '{print $1}' | head -1)
    if [[ -n "$swap_file" ]] && [[ "$swap_file" != "/swapfile" ]]; then
        print_color $YELLOW "⚠️  检测到Swap文件路径: $swap_file"
        print_color $YELLOW "   本脚本默认使用/swapfile，可能需要调整"
        read -p "是否继续使用路径: $swap_file? (y/N): " confirm
        if [[ ! $confirm =~ [Yy] ]]; then
            print_color $GREEN "操作已取消"
            exit 0
        fi
        SWAP_FILE_PATH="$swap_file"
    else
        SWAP_FILE_PATH="/swapfile"
    fi
    
    print_color $GREEN "✓ 系统兼容性检查通过"
    echo ""
}

# 检查磁盘空间
check_disk_space() {
    local required_mb=$1
    local available_mb=$(df -m / | awk 'NR==2 {print $4}')
    local available_gb=$(calculate_gb $available_mb)
    local required_gb=$(calculate_gb $required_mb)
    
    echo "需要空间: ${required_gb}GB (${required_mb}MB)"
    echo "可用空间: ${available_gb}GB (${available_mb}MB)"
    
    if [[ $available_mb -lt $required_mb ]]; then
        print_color $RED "❌ 磁盘空间不足!"
        print_color $RED "   需要: ${required_mb}MB，可用: ${available_mb}MB"
        print_color $RED "   请释放磁盘空间或选择更小的Swap大小"
        return 1
    else
        local free_after_mb=$((available_mb - required_mb))
        local free_after_gb=$(calculate_gb $free_after_mb)
        print_color $GREEN "✓ 磁盘空间充足"
        print_color $CYAN "   调整后剩余空间: ${free_after_gb}GB"
        return 0
    fi
}

# 检查系统负载
check_system_load() {
    print_color $BLUE "=== 系统负载检查 ==="
    
    local load_avg=$(cat /proc/loadavg | awk '{print $1}')
    local cpu_count=$(nproc)
    
    echo "当前负载: $load_avg"
    echo "CPU核心数: $cpu_count"
    
    # 使用awk计算阈值
    local threshold=$(echo "$cpu_count" | awk '{print int($1 * 0.7 + 0.5)}')
    
    # 使用浮点数比较函数
    if float_compare "$load_avg" ">" "$threshold"; then
        print_color $YELLOW "⚠️  系统负载较高，建议在负载较低时执行此操作"
        read -p "是否继续? (y/N): " confirm
        if [[ ! $confirm =~ [Yy] ]]; then
            print_color $GREEN "操作已取消"
            exit 0
        fi
    else
        print_color $GREEN "✓ 系统负载正常"
    fi
    echo ""
}

# 显示当前Swap状态
show_swap_status() {
    print_color $BLUE "=== 当前Swap状态 ==="
    echo "当前Swap使用情况:"
    free -h | grep -E "^(Mem|Swap)"
    echo ""
    
    local swap_info=$(swapon --show 2>/dev/null)
    if [[ -n "$swap_info" ]]; then
        echo "Swap文件详情:"
        echo "$swap_info"
        
        local current_swap_bytes=$(swapon --show --bytes 2>/dev/null | grep -v "NAME" | awk '{print $3}' | head -1)
        if [[ -n "$current_swap_bytes" ]]; then
            CURRENT_SWAP_MB=$((current_swap_bytes / 1024 / 1024))
            CURRENT_SWAP_GB=$(calculate_gb $CURRENT_SWAP_MB)
            echo ""
            print_color $CYAN "当前Swap大小: ${CURRENT_SWAP_GB}GB (${CURRENT_SWAP_MB}MB)"
        fi
    else
        print_color $YELLOW "⚠️  未找到活跃的Swap文件"
    fi
    echo ""
}

# 显示磁盘空间
show_disk_space() {
    print_color $BLUE "=== 磁盘空间状态 ==="
    df -h / | tail -1
    echo ""
}

# 风险提示
show_warnings() {
    print_color $RED "⚠️  ⚠️  ⚠️  重要风险提示 ⚠️  ⚠️  ⚠️"
    echo ""
    echo "1. 调整Swap过程中系统可能会变慢或暂时无响应"
    echo "2. 操作期间如果断电或系统崩溃可能导致数据丢失"
    echo "3. 请确保有足够的磁盘空间"
    echo "4. 建议在系统负载较低时执行此操作"
    echo "5. 操作前请保存所有重要工作"
    echo "6. 确保当前使用的是Swap文件而不是Swap分区"
    echo ""
    print_color $YELLOW "继续操作即表示您已了解并接受上述风险"
    echo ""
}

# 获取用户输入的大小
get_swap_size() {
    print_color $BLUE "=== 选择Swap大小 ==="
    echo ""
    echo "请选择新的Swap大小:"
    echo "1) 512MB   (适合极小内存系统)"
    echo "2) 1GB     (适合轻量使用)"
    echo "3) 1.5GB   (平衡选择)"
    echo "4) 2GB     (推荐用于2GB以下内存)"
    echo "5) 4GB     (标准大小)"
    echo "6) 8GB     (大内存系统备用)"
    echo "7) 自定义大小"
    echo ""
    
    while true; do
        read -p "请输入选项 (1-7): " choice
        case $choice in
            1) SWAP_SIZE_MB=512; break ;;
            2) SWAP_SIZE_MB=1024; break ;;
            3) SWAP_SIZE_MB=1536; break ;;
            4) SWAP_SIZE_MB=2048; break ;;
            5) SWAP_SIZE_MB=4096; break ;;
            6) SWAP_SIZE_MB=8192; break ;;
            7) 
                while true; do
                    read -p "请输入Swap大小 (单位: MB): " custom_size
                    if [[ $custom_size =~ ^[0-9]+$ ]] && [ $custom_size -gt 0 ]; then
                        SWAP_SIZE_MB=$custom_size
                        break
                    else
                        print_color $RED "请输入有效的正整数"
                    fi
                done
                break
                ;;
            *) print_color $RED "无效选项，请重新输入" ;;
        esac
    done
    
    # 计算GB显示
    SWAP_SIZE_GB=$(calculate_gb $SWAP_SIZE_MB)
    echo ""
    print_color $GREEN "您选择的大小: ${SWAP_SIZE_MB}MB (${SWAP_SIZE_GB}GB)"
    
    # 检查磁盘空间
    if ! check_disk_space $SWAP_SIZE_MB; then
        print_color $RED "磁盘空间检查失败，请选择更小的Swap大小或释放磁盘空间"
        return 1
    fi
    echo ""
}

# 确认操作
confirm_operation() {
    echo ""
    print_color $YELLOW "即将执行以下操作:"
    echo "1. 关闭当前Swap (如果存在)"
    echo "2. 删除旧的Swap文件: $SWAP_FILE_PATH"
    echo "3. 创建新的Swap文件: ${SWAP_SIZE_MB}MB"
    echo "4. 设置权限并启用新的Swap"
    echo ""
    
    while true; do
        read -p "是否继续? (y/N): " confirm
        case $confirm in
            [Yy]* ) break;;
            [Nn]* ) 
                print_color $GREEN "操作已取消"
                exit 0
                ;;
            * ) 
                print_color $RED "请输入 y 或 n"
                ;;
        esac
    done
}

# 执行Swap调整
adjust_swap() {
    print_color $BLUE "=== 开始调整Swap ==="
    
    # 关闭当前Swap
    print_color $YELLOW "步骤 1/5: 关闭当前Swap..."
    if sudo swapoff "$SWAP_FILE_PATH" 2>/dev/null; then
        print_color $GREEN "✓ 已关闭当前Swap"
    else
        print_color $CYAN "ℹ️  当前没有活跃的Swap或关闭失败（可能是首次创建）"
    fi
    
    # 删除旧文件
    print_color $YELLOW "步骤 2/5: 删除旧Swap文件..."
    if sudo rm -f "$SWAP_FILE_PATH"; then
        print_color $GREEN "✓ 已删除旧Swap文件"
    else
        print_color $CYAN "ℹ️  旧Swap文件不存在或删除失败"
    fi
    
    # 创建新Swap文件
    print_color $YELLOW "步骤 3/5: 创建新的Swap文件 (${SWAP_SIZE_MB}MB)..."
    if sudo dd if=/dev/zero of="$SWAP_FILE_PATH" bs=1M count=$SWAP_SIZE_MB status=progress; then
        print_color $GREEN "✓ 已创建新Swap文件"
    else
        print_color $RED "❌ 创建Swap文件失败"
        exit 1
    fi
    
    # 设置权限
    print_color $YELLOW "步骤 4/5: 设置文件权限..."
    if sudo chmod 600 "$SWAP_FILE_PATH"; then
        print_color $GREEN "✓ 已设置文件权限"
    else
        print_color $RED "❌ 设置文件权限失败"
        exit 1
    fi
    
    # 格式化为Swap并启用
    print_color $YELLOW "步骤 5/5: 格式化并启用Swap..."
    if sudo mkswap "$SWAP_FILE_PATH" && sudo swapon "$SWAP_FILE_PATH"; then
        print_color $GREEN "✓ 已格式化并启用Swap"
    else
        print_color $RED "❌ 格式化或启用Swap失败"
        exit 1
    fi
    
    print_color $GREEN "✅ Swap调整完成!"
}

# 验证结果
verify_result() {
    print_color $BLUE "=== 验证调整结果 ==="
    echo ""
    echo "新的Swap状态:"
    swapon --show
    echo ""
    free -h | grep -E "^(Mem|Swap)"
    echo ""
    
    # 检查fstab配置
    if grep -q "$SWAP_FILE_PATH" /etc/fstab; then
        print_color $GREEN "✅ /etc/fstab中已配置开机自动挂载"
    else
        print_color $YELLOW "⚠️  /etc/fstab中未配置开机自动挂载，添加中..."
        echo "$SWAP_FILE_PATH none swap sw 0 0" | sudo tee -a /etc/fstab
        print_color $GREEN "✅ 已添加到/etc/fstab"
    fi
}

# 显示使用建议
show_advice() {
    print_color $BLUE "=== 使用建议 ==="
    echo ""
    echo "1. 监控Swap使用情况:"
    echo "   命令: free -h 或 swapon --show"
    echo ""
    echo "2. 如果发现Swap使用率经常接近100%，可以考虑再次增加大小"
    echo ""
    echo "3. 调整内存回收倾向 (可选):"
    echo "   查看当前值: cat /proc/sys/vm/swappiness"
    echo "   临时调整: sudo sysctl vm.swappiness=值 (通常30-60)"
    echo ""
    
    # 如果bc未安装，显示安装建议
    if [[ $HAS_BC -eq 0 ]]; then
        echo ""
        print_color $CYAN "💡 提示: 安装bc可以获得更好的计算精度"
        print_color $CYAN "   运行: sudo apt-get install bc"
    fi
    
    print_color $GREEN "✅ 所有操作已完成!"
}

# 生成使用报告
generate_report() {
    local report_file="/tmp/swap_adjustment_report_$(date +%Y%m%d_%H%M%S).txt"
    
    {
        echo "Swap调整报告 - $(date)"
        echo "================================"
        echo "调整前大小: ${CURRENT_SWAP_GB}GB (${CURRENT_SWAP_MB}MB)"
        echo "调整后大小: ${SWAP_SIZE_GB}GB (${SWAP_SIZE_MB}MB)"
        echo "Swap文件路径: $SWAP_FILE_PATH"
        echo "计算引擎: $([[ $HAS_BC -eq 1 ]] && echo "bc" || echo "awk")"
        echo "操作状态: 成功"
        echo "时间: $(date)"
        echo ""
        echo "当前状态:"
        free -h
        echo ""
        swapon --show
    } > "$report_file"
    
    print_color $CYAN "📊 详细报告已保存至: $report_file"
}

# 主函数
main() {
    print_color $GREEN "======= Swap大小调整脚本（优化依赖版）========"
    echo ""
    
    # 检查依赖
    check_dependencies
    check_root
    check_compatibility
    check_system_load
    show_swap_status
    show_disk_space
    
    # 记录当前Swap大小
    local current_swap_info=$(swapon --show --bytes 2>/dev/null | grep -v "NAME" | head -1)
    if [[ -n "$current_swap_info" ]]; then
        CURRENT_SWAP_MB=$(echo "$current_swap_info" | awk '{print $3/1024/1024}')
        CURRENT_SWAP_GB=$(calculate_gb $CURRENT_SWAP_MB)
    fi
    
    show_warnings
    
    # 获取用户输入并检查磁盘空间
    if ! get_swap_size; then
        exit 1
    fi
    
    confirm_operation
    adjust_swap
    verify_result
    generate_report
    show_advice
}

# 脚本入口
if [[ "${BASH_SOURCE[0]}" == "${0}" ]]; then
    main
fi